"""
RAGFlow文档上传脚本
"""
import os
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
#批量上传文档到ragflow
# === 配置区域 ===
RAGFLOW_BASE_URL = "http://118.....165.74:80/api/v1"  # 例如 "http://localhost:80/api/v1"
API_KEY = "ragflow-UxYmUyNmU....BiNWI0N2EzOW"
KNOWLEDGE_BASE_ID = "425db196ae6....97a39dd136524"
FOLDER_PATH_TO_UPLOAD = "/Users/.../blog/other"  # 要上传的文件夹路径

# 支持的文件类型列表（可根据需要扩展）
SUPPORTED_EXTENSIONS = {
    '.pdf', '.doc', '.docx', '.txt', '.md',
    '.csv', '.xls', '.xlsx',
    '.jpg', '.jpeg', '.png', '.tif', '.gif',
    '.ppt', '.pptx'
}

# === 创建带重试机制的会话 ===
session = requests.Session()
retries = Retry(total=3, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
session.mount("http://", HTTPAdapter(max_retries=retries))
session.mount("https://", HTTPAdapter(max_retries=retries))

headers = {
    "Authorization": f"Bearer {API_KEY}",
    # "Content-Type": 对于文件上传，由requests库自动设置为multipart/form-data
}


def upload_file_to_ragflow(file_path, knowledge_base_id):
    """
    将单个文件上传到指定的RAGFlow知识库 - 使用正确的API端点
    """
    # 使用正确的API端点
    upload_url = f"{RAGFLOW_BASE_URL}/datasets/{knowledge_base_id}/documents"

    try:
        with open(file_path, 'rb') as f:
            files = {'file': (os.path.basename(file_path), f)}
            # 注意：这里不需要额外传递kb_id，因为它已经在URL中

            response = session.post(upload_url, headers=headers, files=files, timeout=60)

        if response.status_code == 200:
            print(f"✅ 上传成功: {file_path}")

            # 尝试获取文档ID并触发解析
            try:
                document_info = response.json()
                document_id = document_info.get('id')
                if document_id:
                    # 触发文档解析
                    trigger_parsing(knowledge_base_id, document_id)
            except Exception as e:
                print(f"⚠️  无法获取文档ID或触发解析: {str(e)}")

            return True
        else:
            print(f"❌ 上传失败 [{response.status_code}]: {file_path} - {response.text}")
            return False

    except Exception as e:
        print(f"💥 上传时发生异常 [{file_path}]: {str(e)}")
        return False


def trigger_parsing(knowledge_base_id, document_id):
    """
    触发文档解析
    """
    try:
        parse_url = f"{RAGFLOW_BASE_URL}/datasets/{knowledge_base_id}/chunks"
        parse_data = {
            "document_ids": [document_id]
        }

        response = session.post(parse_url, headers=headers, json=parse_data, timeout=30)

        if response.status_code == 200:
            print(f"✅ 已触发文档解析: {document_id}")
        else:
            print(f"⚠️  触发文档解析失败: {response.status_code} - {response.text}")

    except Exception as e:
        print(f"💥 触发解析时发生异常: {str(e)}")


def traverse_and_upload(folder_path, knowledge_base_id):
    """
    递归遍历文件夹，上传支持的文件
    """
    # 统计信息
    total_files = 0
    success_files = 0

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in SUPPORTED_EXTENSIONS:
                total_files += 1
                full_file_path = os.path.join(root, file)
                print(f"🔄 正在处理 ({success_files + 1}/{total_files}): {full_file_path}")

                if upload_file_to_ragflow(full_file_path, knowledge_base_id):
                    success_files += 1

                # 添加短暂延迟，避免请求过于频繁
                time.sleep(0.5)
            else:
                print(f"⏭️  跳过不支持的文件类型: {os.path.join(root, file)}")

    print(f"\n📊 上传统计: 成功 {success_files}/{total_files} 个文件")


if __name__ == "__main__":
    import time

    if not os.path.isdir(FOLDER_PATH_TO_UPLOAD):
        print("错误：指定的文件夹路径不存在！")
        exit(1)

    print("开始批量上传文件...")
    traverse_and_upload(FOLDER_PATH_TO_UPLOAD, KNOWLEDGE_BASE_ID)
    print("批量上传任务完成。")